import backtrader as bt

class RsiStrategy(bt.Strategy):
    """
    一个简单的RSI超卖超买策略。
    当RSI指标低于设定的低位阈值（如30）时，买入。
    当RSI指标高于设定的高位阈值（如70）时，卖出。
    """
    params = (
        ('rsi_period', 14),
        ('rsi_overbought', 70),
        ('rsi_oversold', 30),
        ('printlog', True),
    )

    def __init__(self):
        self.dataclose = self.datas[0].close
        self.order = None

        # 计算RSI指标
        self.rsi = bt.indicators.RelativeStrengthIndex(
            period=self.p.rsi_period
        )

    def log(self, txt, dt=None, doprint=False):
        """ 日志记录函数 """
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            print(f'{dt.isoformat()}, {txt}')

    def notify_order(self, order):
        """ 订单状态通知 """
        if order.status in [order.Submitted, order.Accepted]:
            return

        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(f'BUY EXECUTED, Price: {order.executed.price:.2f}, Cost: {order.executed.value:.2f}, Comm {order.executed.comm:.2f}')
            else:
                self.log(f'SELL EXECUTED, Price: {order.executed.price:.2f}, Cost: {order.executed.value:.2f}, Comm {order.executed.comm:.2f}')
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')

        self.order = None

    def next(self):
        """ 策略主逻辑 """
        if self.order:
            return

        if not self.position:
            # 如果没有持仓，并且RSI低于超卖线
            if self.rsi < self.p.rsi_oversold:
                self.log(f'BUY CREATE, RSI: {self.rsi[0]:.2f}, Close: {self.dataclose[0]:.2f}')
                self.order = self.buy()
        else:
            # 如果有持仓，并且RSI高于超买线
            if self.rsi > self.p.rsi_overbought:
                self.log(f'SELL CREATE, RSI: {self.rsi[0]:.2f}, Close: {self.dataclose[0]:.2f}')
                self.order = self.sell() 